#include <sys/types.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <syslog.h>
#include <fcntl.h>

#ifndef MIN
#define MIN(a,b) ((a) < (b) ? (a) : (b))
#endif

#include "fdwatch.h"

#ifdef HAVE_SELECT
#ifndef FD_SET
#define NFDBITS         32
#define FD_SETSIZE      32
#define FD_SET(n, p)    ((p)->fds_bits[(n)/NFDBITS] |= (1 << ((n) % NFDBITS)))
#define FD_CLR(n, p)    ((p)->fds_bits[(n)/NFDBITS] &= ~(1 << ((n) % NFDBITS)))
#define FD_ISSET(n, p)  ((p)->fds_bits[(n)/NFDBITS] & (1 << ((n) % NFDBITS)))
#define FD_ZERO(p)      bzero((char*)(p), sizeof(*(p)))
#endif /* !FD_SET */
#endif /* HAVE_SELECT */

static int nfiles;
static long nwatches;
static int* fd_rw;
static void** fd_data;
static int nreturned, next_ridx;

#ifdef HAVE_SELECT

#define WHICH                  "select"
#define INIT( nfiles )         select_init( nfiles )
#define ADD_FD( fd, rw )       select_add_fd( fd, rw )
#define DEL_FD( fd )           select_del_fd( fd )
#define WATCH( timeout_msecs ) select_watch( timeout_msecs )
#define CHECK_FD( fd )         select_check_fd( fd )
#define GET_FD( ridx )         select_get_fd( ridx )

static int select_init( int nfiles );
static void select_add_fd( int fd, int rw );
static void select_del_fd( int fd );
static int select_watch( long timeout_msecs );
static int select_check_fd( int fd );
static int select_get_fd( int ridx );

#endif /* HAVE_SELECT */

/* Routines. */

/* Figure out how many file descriptors the system allows, and
** initialize the fdwatch data structures.  Returns -1 on failure.
*/
int fdwatch_get_nfiles( void )
{
    int i;
#ifdef RLIMIT_NOFILE
    struct rlimit rl;
#endif /* RLIMIT_NOFILE */

    /* Figure out how many fd's we can have. */
    nfiles = getdtablesize();
#ifdef RLIMIT_NOFILE
    /* If we have getrlimit(), use that, and attempt to raise the limit. */
    if ( getrlimit( RLIMIT_NOFILE, &rl ) == 0 )
    {
        nfiles = rl.rlim_cur;
        if ( rl.rlim_max == RLIM_INFINITY )
            rl.rlim_cur = 8192;         /* arbitrary */
        else if ( rl.rlim_max > rl.rlim_cur )
            rl.rlim_cur = rl.rlim_max;
        if ( setrlimit( RLIMIT_NOFILE, &rl ) == 0 )
            nfiles = rl.rlim_cur;
    }
#endif /* RLIMIT_NOFILE */

#if defined(HAVE_SELECT) && ! ( defined(HAVE_POLL) || defined(HAVE_DEVPOLL) || defined(HAVE_KQUEUE) )
    /* If we use select(), then we must limit ourselves to FD_SETSIZE. */
    nfiles = MIN( nfiles, FD_SETSIZE );
#endif /* HAVE_SELECT && ! ( HAVE_POLL || HAVE_DEVPOLL || HAVE_KQUEUE ) */

    /* Initialize the fdwatch data structures. */
    nwatches = 0;
    fd_rw = (int*) malloc( sizeof(int) * nfiles );
    fd_data = (void**) malloc( sizeof(void*) * nfiles );
    if ( fd_rw == (int*) 0 || fd_data == (void**) 0 )
        return -1;
    for ( i = 0; i < nfiles; ++i )
        fd_rw[i] = -1;
    if ( INIT( nfiles ) == -1 )
        return -1;

    return nfiles;
}


/* Add a descriptor to the watch list.  rw is either FDW_READ or FDW_WRITE.  */
void fdwatch_add_fd( int fd, void* client_data, int rw )
{
    if ( fd < 0 || fd >= nfiles || fd_rw[fd] != -1 )
    {
        syslog( LOG_ERR, "bad fd (%d) passed to fdwatch_add_fd!", fd );
        return;
    }
    ADD_FD( fd, rw );
    fd_rw[fd] = rw;
    fd_data[fd] = client_data;
}


/* Remove a descriptor from the watch list. */
void fdwatch_del_fd( int fd )
{
    if ( fd < 0 || fd >= nfiles || fd_rw[fd] == -1 )
    {
        syslog( LOG_ERR, "bad fd (%d) passed to fdwatch_del_fd!", fd );
        return;
    }
    DEL_FD( fd );
    fd_rw[fd] = -1;
    fd_data[fd] = (void*) 0;
}

/* Do the watch.  Return value is the number of descriptors that are ready,
** or 0 if the timeout expired, or -1 on errors.  A timeout of INFTIM means
** wait indefinitely.
*/
int fdwatch( long timeout_msecs )
{
    ++nwatches;
    nreturned = WATCH( timeout_msecs );
    next_ridx = 0;
    return nreturned;
}


/* Check if a descriptor was ready. */
int fdwatch_check_fd( int fd )
{
    if ( fd < 0 || fd >= nfiles || fd_rw[fd] == -1 )
    {
        syslog( LOG_ERR, "bad fd (%d) passed to fdwatch_check_fd!", fd );
        return 0;
    }
    return CHECK_FD( fd );
}


void* fdwatch_get_next_client_data( void )
{
    int fd;

    if ( next_ridx >= nreturned )
        return (void*) -1;
    fd = GET_FD( next_ridx++ );
    if ( fd < 0 || fd >= nfiles )
        return (void*) 0;
    return fd_data[fd];
}


/* Generate debugging statistics syslog message. */
void fdwatch_logstats( long secs )
{
    if ( secs > 0 )
        syslog(
                LOG_INFO, "  fdwatch - %ld %ss (%g/sec)",
                nwatches, WHICH, (float) nwatches / secs );
    nwatches = 0;
}



#ifdef HAVE_SELECT

static fd_set master_rfdset;
static fd_set master_wfdset;
static fd_set working_rfdset;
static fd_set working_wfdset;
static int* select_fds;
static int* select_fdidx;
static int* select_rfdidx;
static int nselect_fds;
static int maxfd;
static int maxfd_changed;


static int select_init( int nfiles )
{
    int i;

    FD_ZERO( &master_rfdset );
    FD_ZERO( &master_wfdset );
    select_fds = (int*) malloc( sizeof(int) * nfiles );
    select_fdidx = (int*) malloc( sizeof(int) * nfiles );
    select_rfdidx = (int*) malloc( sizeof(int) * nfiles );
    if ( select_fds == (int*) 0 || select_fdidx == (int*) 0 ||
            select_rfdidx == (int*) 0 )
        return -1;
    nselect_fds = 0;
    maxfd = -1;
    maxfd_changed = 0;
    for ( i = 0; i < nfiles; ++i )
        select_fds[i] = select_fdidx[i] = -1;
    return 0;
}


static void select_add_fd( int fd, int rw )
{
    if ( nselect_fds >= nfiles )
    {
        syslog( LOG_ERR, "too many fds in select_add_fd!" );
        return;
    }
    select_fds[nselect_fds] = fd;
    switch ( rw )
    {
        case FDW_READ: FD_SET( fd, &master_rfdset ); break;
        case FDW_WRITE: FD_SET( fd, &master_wfdset ); break;
        default: break;
    }
    if ( fd > maxfd )
        maxfd = fd;
    select_fdidx[fd] = nselect_fds;
    ++nselect_fds;
}


static void select_del_fd( int fd )
{
    int idx = select_fdidx[fd];

    if ( idx < 0 || idx >= nfiles )
    {
        syslog( LOG_ERR, "bad idx (%d) in select_del_fd!", idx );
        return;
    }

    --nselect_fds;
    select_fds[idx] = select_fds[nselect_fds];
    select_fdidx[select_fds[idx]] = idx;
    select_fds[nselect_fds] = -1;
    select_fdidx[fd] = -1;

    FD_CLR( fd, &master_rfdset );
    FD_CLR( fd, &master_wfdset );

    if ( fd >= maxfd )
        maxfd_changed = 1;
}


static int select_get_maxfd( void )
{
    if ( maxfd_changed )
    {
        int i;
        maxfd = -1;
        for ( i = 0; i < nselect_fds; ++i )
            if ( select_fds[i] > maxfd )
                maxfd = select_fds[i];
        maxfd_changed = 0;
    }
    return maxfd;
}


static int select_watch( long timeout_msecs )
{
    int mfd;
    int r, idx, ridx;

    working_rfdset = master_rfdset;
    working_wfdset = master_wfdset;
    mfd = select_get_maxfd();
    if ( timeout_msecs == INFTIM )
        r = select(
                mfd + 1, &working_rfdset, &working_wfdset, (fd_set*) 0,
                (struct timeval*) 0 );
    else
    {
        struct timeval timeout;
        timeout.tv_sec = timeout_msecs / 1000L;
        timeout.tv_usec = ( timeout_msecs % 1000L ) * 1000L;
        r = select(
                mfd + 1, &working_rfdset, &working_wfdset, (fd_set*) 0, &timeout );
    }
    if ( r <= 0 )
        return r;

    ridx = 0;
    for ( idx = 0; idx < nselect_fds; ++idx )
        if ( select_check_fd( select_fds[idx] ) )
        {
            select_rfdidx[ridx++] = select_fds[idx];
            if ( ridx == r )
                break;
        }

    return ridx;	/* should be equal to r */
}


static int select_check_fd( int fd )
{
    switch ( fd_rw[fd] )
    {
        case FDW_READ: return FD_ISSET( fd, &working_rfdset );
        case FDW_WRITE: return FD_ISSET( fd, &working_wfdset );
        default: return 0;
    }
}


static int select_get_fd( int ridx )
{
    if ( ridx < 0 || ridx >= nfiles )
    {
        syslog( LOG_ERR, "bad ridx (%d) in select_get_fd!", ridx );
        return -1;
    }
    return select_rfdidx[ridx];
}

#endif /* HAVE_SELECT */
